import tensorflow as tf
from tensorflow import keras
from sklearn.model_selection import train_test_split
import re
import numpy as np
import time
from data_util import  TextData
from args import Args
from glove import Glove


class Encoder(keras.Model):
    def __init__(self, vocab_size, embedding_dim, enc_units, batch_sz):
        super(Encoder, self).__init__()
        self.batch_sz = batch_sz
        self.enc_units = enc_units
        self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
        self.gru = tf.keras.layers.GRU(self.enc_units,
                                       return_sequences=True,
                                       return_state=True,
                                       recurrent_initializer='glorot_uniform')
    
    def call(self, x, hidden):
        x = self.embedding(x)
        output, state = self.gru(x, initial_state = hidden)
        return output, state
    
    def initialize_hidden_state(self):
        return tf.zeros((self.batch_sz, self.enc_units))

class BahdanauAttention(keras.layers.Layer):
    def __init__(self, units):
        super(BahdanauAttention, self).__init__()
        self.W1 = tf.keras.layers.Dense(units)
        self.W2 = tf.keras.layers.Dense(units)
        self.V = tf.keras.layers.Dense(1)
    
    def call(self, query, values):
        # 隐藏层的形状 == （批大小，隐藏层大小）
        # hidden_with_time_axis 的形状 == （批大小，1，隐藏层大小）
        # 这样做是为了执行加法以计算分数  
        hidden_with_time_axis = tf.expand_dims(query, 1)
        
        # 分数的形状 == （批大小，最大长度，1）
        # 我们在最后一个轴上得到 1， 因为我们把分数应用于 self.V
        # 在应用 self.V 之前，张量的形状是（批大小，最大长度，单位）
        score = self.V(tf.nn.tanh(
            self.W1(values) + self.W2(hidden_with_time_axis)))
        
        # 注意力权重 （attention_weights） 的形状 == （批大小，最大长度，1）
        attention_weights = tf.nn.softmax(score, axis=1)
        
        # 上下文向量 （context_vector） 求和之后的形状 == （批大小，隐藏层大小）
        context_vector = attention_weights * values
        context_vector = tf.reduce_sum(context_vector, axis=1)
        
        return context_vector, attention_weights

class Decoder(keras.Model):
    def __init__(self, vocab_size, embedding_dim, dec_units, batch_sz):
        super(Decoder, self).__init__()
        self.batch_sz = batch_sz
        self.dec_units = dec_units
        self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
        self.gru = tf.keras.layers.GRU(self.dec_units,
                                       return_sequences=True,
                                       return_state=True,
                                       recurrent_initializer='glorot_uniform')
        self.fc = tf.keras.layers.Dense(vocab_size)
        
        # 用于注意力
        self.attention = BahdanauAttention(self.dec_units)
    
    def call(self, x, hidden, enc_output):
        # 编码器输出 （enc_output） 的形状 == （批大小，最大长度，隐藏层大小）
        context_vector, attention_weights = self.attention(hidden, enc_output)
        
        # x 在通过嵌入层后的形状 == （批大小，1，嵌入维度）
        x = self.embedding(x)
        
        # x 在拼接 （concatenation） 后的形状 == （批大小，1，嵌入维度 + 隐藏层大小）
        x = tf.concat([tf.expand_dims(context_vector, 1), x], axis=-1)
        
        # 将合并后的向量传送到 GRU
        output, state = self.gru(x)
        
        # 输出的形状 == （批大小 * 1，隐藏层大小）
        output = tf.reshape(output, (-1, output.shape[2]))
        
        # 输出的形状 == （批大小，vocab）
        x = self.fc(output)
        
        return x, state, attention_weights

def loss_function(real, pred):
    mask = tf.math.logical_not(tf.math.equal(real, 0))
    loss_ = loss_object(real, pred)
    
    mask = tf.cast(mask, dtype=loss_.dtype)
    loss_ *= mask
    
    return tf.reduce_mean(loss_)

@tf.function
def train_step(inp, targ, enc_hidden):
    loss = 0
    
    with tf.GradientTape() as tape:
        enc_output, enc_hidden = encoder(inp, enc_hidden)
        
        dec_hidden = enc_hidden
        
        dec_input = tf.expand_dims([1] * BATCH_SIZE, 1)
        
        # 教师强制 - 将目标词作为下一个输入
        for t in range(1, targ.shape[1]):
          # 将编码器输出 （enc_output） 传送至解码器
          predictions, dec_hidden, _ = decoder(dec_input, dec_hidden, enc_output)
        
          loss += loss_function(targ[:, t], predictions)
        
        
          # 使用教师强制
          dec_input = tf.expand_dims(targ[:, t], 1)
    
    batch_loss = (loss / int(targ.shape[1]))
    
    variables = encoder.trainable_variables + decoder.trainable_variables
    
    gradients = tape.gradient(loss, variables)
    
    optimizer.apply_gradients(zip(gradients, variables))
    
    return batch_loss


if __name__ == '__main__':

    args = Args()    
    textData = TextData(args)
    
    inp = np.array(textData.train_samples)[:,0]
    targ = np.array(textData.train_samples)[:,1]
    for i in range(len(inp)):
        inp[i] = [1] + inp[i] + [2]
        targ[i] = [1] + targ[i] + [2]
    
    input_tensor = tf.keras.preprocessing.sequence.pad_sequences(inp, padding='post')
    target_tensor = tf.keras.preprocessing.sequence.pad_sequences(targ, padding='post')
    
    max_length = 18
    
    input_tensor_train, input_tensor_val, target_tensor_train, target_tensor_val = train_test_split(input_tensor, target_tensor, test_size=0.1)
    
    BUFFER_SIZE = len(input_tensor_train)
    BATCH_SIZE = 512
    steps_per_epoch = len(input_tensor)//BATCH_SIZE
    embedding_dim = 64
    units = 512
    vocab_size = textData.vocab_size

    
    dataset = tf.data.Dataset.from_tensor_slices((input_tensor_train, target_tensor_train)).shuffle(BUFFER_SIZE)
    dataset = dataset.batch(BATCH_SIZE, drop_remainder=True)
    
    encoder = Encoder(vocab_size, embedding_dim, units, BATCH_SIZE)
    decoder = Decoder(vocab_size, embedding_dim, units, BATCH_SIZE)
    
    #载入glove词向量作为encoder和decoder嵌入层预训练
    glove = Glove(vocab_size, embedding_dim)
    checkpoint = tf.train.Checkpoint(glove=glove)
    checkpoint.restore(tf.train.latest_checkpoint("./save_model/glove"))
    embedding_weights = glove.get_uaddv_weights()
    encoder.embedding(tf.zeros(0,dtype="int32"))
    decoder.embedding(tf.zeros(0,dtype="int32"))
    encoder.embedding.set_weights(embedding_weights) 
    decoder.embedding.set_weights(embedding_weights) 
    
    optimizer = tf.keras.optimizers.Adam()
    loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
        from_logits=True, reduction='none')
    
    #保存模型
    checkpoint = tf.train.Checkpoint(optimizer=optimizer,
                                     encoder=encoder,
                                     decoder=decoder)
    #tensorboard数据可视化
    log_dir = 'logs/seq2seq_log'
    summary_writer = tf.summary.create_file_writer(log_dir)

    EPOCHS = 100
    for epoch in range(EPOCHS):
        start = time.time()
        
        enc_hidden = encoder.initialize_hidden_state()
        total_loss = 0
        
        for (batch, (inp, targ)) in enumerate(dataset):
            batch_loss = train_step(inp, targ, enc_hidden)
            total_loss += batch_loss
            
            if batch % 100 == 0:
                print('Epoch {} Batch {} Loss {:.4f}'.format(epoch + 1,
                                                             batch,
                                                             batch_loss.numpy()))
        #记录每轮损失
        with summary_writer.as_default():
            tf.summary.scalar('loss_mse', data=total_loss / steps_per_epoch ,step=epoch)
        
        print('Epoch {} Loss {:.4f}'.format(epoch + 1,
                                          total_loss / steps_per_epoch))
        print('Time taken for 1 epoch {} sec\n'.format(time.time() - start))
        
    checkpoint.save('./save_model/seq2seq/seq2seq.ckpt')
        